from taa_core.modeling.roi_heads.box_head import loss
import numpy as np
import torch
from torch.nn import functional as F
from torch import device, dtype, nn
import os
from ..utils import concat_box_prediction_layers
from taa_core.layers import smooth_l1_loss
from taa_core.layers import SigmoidFocalLoss
from taa_core.modeling.matcher import Matcher
from taa_core.structures.boxlist_ops import boxlist_iou
from taa_core.structures.boxlist_ops import cat_boxlist
from .sample_prob import SampleProbNet
from .utils import concat_cls_reg_feats


INF = 100000000


def get_num_gpus():
    return int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1


def reduce_sum(tensor):
    if get_num_gpus() <= 1:
        return tensor
    import torch.distributed as dist
    tensor = tensor.clone()
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    return tensor


class TAALossComputation(nn.Module):
    def __init__(self, cfg, box_coder):
        super(TAALossComputation, self).__init__()

        self.cfg = cfg
        self.cls_loss_func = SigmoidFocalLoss(cfg.MODEL.PAA.LOSS_GAMMA,
                                              cfg.MODEL.PAA.LOSS_ALPHA)
        self.iou_pred_loss_func = nn.BCEWithLogitsLoss(reduction="sum")
        self.matcher = Matcher(cfg.MODEL.PAA.IOU_THRESHOLD,
                               cfg.MODEL.PAA.IOU_THRESHOLD,
                               True)
        self.box_coder = box_coder
        self.fpn_strides=[8, 16, 32, 64, 128]
        self.reg_loss_type = cfg.MODEL.PAA.REG_LOSS_TYPE
        
        self.sync_num_pos = cfg.MODEL.TAA.USE_SYNC_POS
        self.num_classes = cfg.MODEL.TAA.NUM_CLASSES - 1
        self.sample_prob_net = SampleProbNet()

    def GIoULoss(self, pred, target, anchor, weight=None):
        pred_boxes = self.box_coder.decode(pred.view(-1, 4), anchor.view(-1, 4))
        pred_x1 = pred_boxes[:, 0]
        pred_y1 = pred_boxes[:, 1]
        pred_x2 = pred_boxes[:, 2]
        pred_y2 = pred_boxes[:, 3]
        pred_x2 = torch.max(pred_x1, pred_x2)
        pred_y2 = torch.max(pred_y1, pred_y2)
        pred_area = (pred_x2 - pred_x1) * (pred_y2 - pred_y1)

        gt_boxes = self.box_coder.decode(target.view(-1, 4), anchor.view(-1, 4))
        target_x1 = gt_boxes[:, 0]
        target_y1 = gt_boxes[:, 1]
        target_x2 = gt_boxes[:, 2]
        target_y2 = gt_boxes[:, 3]
        target_area = (target_x2 - target_x1) * (target_y2 - target_y1)

        x1_intersect = torch.max(pred_x1, target_x1)
        y1_intersect = torch.max(pred_y1, target_y1)
        x2_intersect = torch.min(pred_x2, target_x2)
        y2_intersect = torch.min(pred_y2, target_y2)
        area_intersect = torch.zeros(pred_x1.size()).to(pred)
        mask = (y2_intersect > y1_intersect) * (x2_intersect > x1_intersect)
        area_intersect[mask] = (x2_intersect[mask] - x1_intersect[mask]) * (y2_intersect[mask] - y1_intersect[mask])

        x1_enclosing = torch.min(pred_x1, target_x1)
        y1_enclosing = torch.min(pred_y1, target_y1)
        x2_enclosing = torch.max(pred_x2, target_x2)
        y2_enclosing = torch.max(pred_y2, target_y2)
        area_enclosing = (x2_enclosing - x1_enclosing) * (y2_enclosing - y1_enclosing) + 1e-7

        area_union = pred_area + target_area - area_intersect + 1e-7
        ious = area_intersect / area_union
        gious = ious - (area_enclosing - area_union) / area_enclosing

        pos_losses = 1 - gious
        neg_losses = gious

        if weight is not None and weight.sum() > 0:
            return (pos_losses * weight), (neg_losses * weight)
        else:
            assert pos_losses.numel() != 0 and neg_losses.numel() != 0
            return pos_losses, neg_losses

    def prepare_iou_based_targets(self, targets, anchors):
        """Compute IoU-based targets"""

        cls_labels = []
        reg_targets = []
        matched_idx_all = []
        for im_i in range(len(targets)):
            targets_per_im = targets[im_i]
            assert targets_per_im.mode == "xyxy"
            bboxes_per_im = targets_per_im.bbox
            labels_per_im = targets_per_im.get_field("labels")
            anchors_per_im = cat_boxlist(anchors[im_i])
            num_gt = bboxes_per_im.shape[0]

            match_quality_matrix = boxlist_iou(targets_per_im, anchors_per_im)
            matched_idxs = self.matcher(match_quality_matrix)
            targets_per_im = targets_per_im.copy_with_fields(['labels'])
            matched_targets = targets_per_im[matched_idxs.clamp(min=0)]

            cls_labels_per_im = matched_targets.get_field("labels")
            cls_labels_per_im = cls_labels_per_im.to(dtype=torch.float32)

            # Background (negative examples)
            bg_indices = matched_idxs == Matcher.BELOW_LOW_THRESHOLD
            cls_labels_per_im[bg_indices] = 0

            # discard indices that are between thresholds
            inds_to_discard = matched_idxs == Matcher.BETWEEN_THRESHOLDS
            cls_labels_per_im[inds_to_discard] = -1

            matched_gts = matched_targets.bbox
            matched_idx_all.append(matched_idxs.view(1, -1))

            reg_targets_per_im = self.box_coder.encode(matched_gts, anchors_per_im.bbox)
            cls_labels.append(cls_labels_per_im)
            reg_targets.append(reg_targets_per_im)

        return cls_labels, reg_targets, matched_idx_all

    def compute_reg_loss(
        self, regression_targets, box_regression, all_anchors, labels, weights
    ):
        if 'iou' in self.reg_loss_type:
            reg_pos_loss, reg_neg_loss = self.GIoULoss(box_regression,
                                     regression_targets,
                                     all_anchors,
                                     weight=weights)
        elif self.reg_loss_type == 'smoothl1':
            raise NotImplementedError
            # reg_loss = smooth_l1_loss(box_regression,
            #                           regression_targets,
            #                           beta=self.bbox_reg_beta,
            #                           size_average=False,
            #                           sum=False)
            # if weights is not None:
            #     reg_loss *= weights
        else:
            raise NotImplementedError
        return reg_pos_loss[labels > 0].view(-1), reg_neg_loss[labels > 0].view(-1)

    def compute_ious(self, boxes1, boxes2):
        area1 = (boxes1[:, 2] - boxes1[:, 0] + 1) * (boxes1[:, 3] - boxes1[:, 1] + 1)
        area2 = (boxes2[:, 2] - boxes2[:, 0] + 1) * (boxes2[:, 3] - boxes2[:, 1] + 1)
        lt = torch.max(boxes1[:, :2], boxes2[:, :2])
        rb = torch.min(boxes1[:, 2:], boxes2[:, 2:])
        wh = (rb - lt + 1).clamp(min=0)
        inter = wh[:, 0] * wh[:, 1]
        return inter / (area1 + area2 - inter)

    def forward(self, box_cls, box_regression, targets, anchors, locations):
        device = box_cls[0].device
        # get IoU-based anchor assignment first to compute anchor scores
        (iou_based_labels,
         iou_based_reg_targets,
         matched_idx_all) = self.prepare_iou_based_targets(targets, anchors)
        matched_idx_all = torch.cat(matched_idx_all, dim=0)

        N = len(iou_based_labels)
        iou_based_labels_flatten = torch.cat(iou_based_labels, dim=0).int()
        iou_based_reg_targets_flatten = torch.cat(iou_based_reg_targets, dim=0)
        box_cls_flatten, box_regression_flatten = concat_box_prediction_layers(
            box_cls, box_regression)
        anchors_flatten = torch.cat([cat_boxlist(anchors_per_image).bbox
            for anchors_per_image in anchors], dim=0)

        pos_inds = torch.nonzero(iou_based_labels_flatten > 0).squeeze(1)

        num_pos = len(pos_inds)
        num_gpus = get_num_gpus()
        num_pos_avg_per_gpu = max(float(num_pos), 1.0)
        # if self.sync_num_pos:
        #     pos_inds_t = pos_inds.clone()
        #     total_num_pos = reduce_sum(
        #         pos_inds_t.new_tensor([pos_inds_t.numel()])).item()
        #     num_pos_avg_per_gpu = max(total_num_pos / float(num_gpus), 1.0)
        # else:
        #     num_pos_avg_per_gpu = max(float(num_pos), 1.0)

        if pos_inds.numel() > 0:
            with torch.no_grad():
                cls_loss = self.cls_loss_func(box_cls_flatten.detach(),
                                              iou_based_labels_flatten,
                                              sum=False) 
                cls_loss = cls_loss.sum(dim=1)[pos_inds]
                clses = box_cls_flatten[pos_inds]
                gt_boxes = self.box_coder.decode(iou_based_reg_targets_flatten.view(-1, 4), anchors_flatten)
                boxes = self.box_coder.decode(box_regression_flatten, anchors_flatten)
                ious = self.compute_ious(gt_boxes, boxes)
                # class_range = torch.arange(1, self.num_classes+1, dtype=torch.long, device=device).unsqueeze(0)

            anchor_probs = self.sample_prob_net(
                ious=ious[pos_inds],
                probs=clses,
                cls_loss=cls_loss
            )
            
            if self.sync_num_pos:
                pos_inds_t = pos_inds.clone()
                total_num_pos = reduce_sum(
                    # pos_inds_t.new_tensor([pos_inds_t.numel()])).item()
                    ious.sum()).item()
                num_pos_avg_per_gpu = total_num_pos / float(num_gpus)
            else:
                num_pos_avg_per_gpu = float(ious.sum())

            reg_weights = torch.zeros(box_regression_flatten.size(0), device=device).detach()
            # _cls_inds = (iou_based_labels_flatten.unsqueeze(1) == class_range)
            # print(box_cls_flatten[pos_inds, iou_based_labels_flatten[pos_inds].long()-1].shape);exit()
            reg_weights[pos_inds] = torch.sigmoid(box_cls_flatten[pos_inds, iou_based_labels_flatten[pos_inds].long()-1])
            reg_pos_loss, reg_neg_loss = self.compute_reg_loss(iou_based_reg_targets_flatten, 
                                             box_regression_flatten,
                                             anchors_flatten,
                                             iou_based_labels_flatten,
                                             weights=reg_weights)
            loc_pos_prob = (-reg_pos_loss).exp()
            loc_neg_prob = (-reg_neg_loss).exp()

            loss = join_focal_loss(
                box_cls_flatten,
                iou_based_labels_flatten,
                (loc_pos_prob, loc_neg_prob),
                anchor_probs,
                # weight=ious[pos_inds],
                gamma=self.cfg.MODEL.TAA.LOSS_GAMMA,
                alpha=self.cfg.MODEL.TAA.LOSS_ALPHA,
                # test=torch.sigmoid(box_cls_flatten[pos_inds, iou_based_labels_flatten[pos_inds].long()-1])
            ) / num_pos_avg_per_gpu
            
            # return loss
        else:
            loss = self.cls_loss_func(box_cls_flatten, iou_based_labels_flatten) / num_pos_avg_per_gpu
            # reg_loss = box_regression_flatten.sum()
        return loss


def join_focal_loss(logits, labels, loc_probs, sample_probs, weight=None, gamma=2, alpha=0.25, test=None):
    num_classes = logits.size(1)
    dtype = labels.dtype
    device = labels.device
    loc_pos_probs, loc_neg_probs = loc_probs
    pos_probs = sample_probs
    neg_probs = -torch.log(pos_probs)

    logits = torch.sigmoid(logits)
    class_range = torch.arange(1, num_classes+1, dtype=dtype, device=device).unsqueeze(0)
    t = labels.unsqueeze(1)
    pos = logits[t == class_range]
    neg = logits[(t != class_range) * (t >= 0)]

    if weight is not None:
        assert False
        term0 = (1 - weight)*(pos_probs*(pos*loc_pos_probs)).log() + \
                weight*(neg_probs*((1-pos)*loc_neg_probs)).log()
    else:
        term0 = (pos_probs * (pos * loc_pos_probs)).log() + (neg_probs * ((1-pos) * loc_neg_probs)).log()
    
    term1 = neg ** gamma * (1 - neg).log()
    
    print(term0.sum())
    assert torch.isnan(term0).sum()==0, "NaN in term0"
    assert torch.isnan(term1).sum()==0, "NaN in term1"

    return -term0.sum() - term1.sum()


def make_taa_loss_evaluator(cfg, box_coder):
    loss_evaluator = TAALossComputation(cfg, box_coder)
    return loss_evaluator
